#!/usr/bin/env python3
"""
Main entry point for AE-PINN (Attention-Enhanced Physics-Informed Neural Networks).

This script provides a command-line interface for running AE-PINN simulations
on elliptic interface problems. It supports various problem configurations,
model architectures, and training strategies.

Usage:
    python main.py --experiment star_shaped --epochs 10000 --save_dir results/
    python main.py --config config.json
    python main.py --help

Example configurations and experiments are provided for:
- Star-shaped interface problems
- Circular interface problems  
- Custom problem setups
"""

import argparse
import json
import os
import sys
import time
from pathlib import Path
import torch
import numpy as np
import matplotlib
matplotlib.use('Agg')  # Use non-interactive backend for headless plotting

# Add current directory to path for imports
sys.path.append(str(Path(__file__).parent))


def initialize_cuda_context():
    """Initialize CUDA context to avoid cuBLAS warnings."""
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        # Create a small tensor operation to initialize CUDA context
        device = torch.device('cuda')
        dummy = torch.zeros(1, device=device)
        _ = dummy + 1
        torch.cuda.synchronize()
        del dummy

from networks.fcnn import FCNN, create_fcnn
from networks.ia_nn import IANN, CompositeAEPINN, create_star_shaped_level_set
from pinn.loss import AEPINNLoss
from pinn.solver import AEPINNSolver
from experiments.star_shaped_interface import StarShapedExperiment


def create_circular_interface_problem(r0=0.5, alpha_in=1.0, alpha_out=2.0):
    """
    Create a simple circular interface problem for testing.
    
    Args:
        r0 (float): Interface radius
        alpha_in (float): Diffusion coefficient inside
        alpha_out (float): Diffusion coefficient outside
        
    Returns:
        dict: Problem configuration
    """
    def level_set_func(x, y):
        return torch.sqrt(x**2 + y**2) - r0
    
    def alpha_func(x, y):
        phi = level_set_func(x, y)
        return torch.where(phi < 0, alpha_in, alpha_out)
    
    def f_func(x, y):
        return torch.ones_like(x)
    
    def boundary_func(x, y):
        return torch.zeros_like(x)
    
    def jump_func(x, y):
        return 0.1 * torch.sin(2 * np.pi * x) * torch.sin(2 * np.pi * y)
    
    def flux_jump_func(x, y):
        return torch.zeros_like(x)
    
    return {
        'level_set_func': level_set_func,
        'alpha_func': alpha_func,
        'f_func': f_func,
        'boundary_func': boundary_func,
        'jump_func': jump_func,
        'flux_jump_func': flux_jump_func,
        'domain_bounds': [[-1.0, 1.0], [-1.0, 1.0]]
    }


def create_custom_problem_from_config(problem_config):
    """
    Create a custom problem from configuration.
    
    Args:
        problem_config (dict): Problem configuration
        
    Returns:
        dict: Problem functions and parameters
    """
    # This is a placeholder for creating custom problems from config
    # In practice, you would implement specific problem types here
    return create_circular_interface_problem()


def run_star_shaped_experiment(args):
    """Run the star-shaped interface experiment."""
    print("Running star-shaped interface experiment...")
    
    # Configuration
    problem_config = {
        'r0': args.r0,
        'alpha_in': args.alpha_in,
        'alpha_out': args.alpha_out,
        'betas': [0.3, 0.2, 0.1],
        'etas': [3, 4, 5],
        'thetas': [0, np.pi/4, np.pi/2]
    }
    
    model_config = {
        'fcnn_layers': [50, 50, 50],
        'ia_hidden_dim': 40,
        'ia_attention_layers': 3,
        'activation': 'tanh'
    }
    
    training_config = {
        'epochs': args.epochs,
        'learning_rate': args.learning_rate,
        'n_interior': args.n_interior,
        'n_boundary': args.n_boundary,
        'n_interface': args.n_interface,
        'optimizer': args.optimizer,
        'scheduler': 'step',
        'loss_weights': {
            'tau_pde': 1.0,
            'tau_bc': 100.0,
            'tau_interface_u': 100.0,
            'tau_interface_flux': 100.0
        }
    }
    
    # Create and run experiment
    experiment = StarShapedExperiment(
        problem_config=problem_config,
        model_config=model_config,
        training_config=training_config
    )
    
    experiment.run_complete_experiment(save_dir=args.save_dir)


def run_circular_experiment(args):
    """Run a circular interface experiment."""
    print("Running circular interface experiment...")
    
    # Create problem
    problem = create_circular_interface_problem(
        r0=args.r0, 
        alpha_in=args.alpha_in, 
        alpha_out=args.alpha_out
    )
    
    # Create model
    fcnn = FCNN(
        input_dim=2,
        output_dim=1,
        hidden_layers=[50, 50, 50],
        activation='tanh'
    )
    
    ia_nn = IANN(
        input_dim=2,
        hidden_dim=40,
        num_attention_layers=3,
        output_dim=1,
        level_set_func=problem['level_set_func'],
        activation='tanh'
    )
    
    model = CompositeAEPINN(fcnn, ia_nn)
    
    # Create loss function
    loss_function = AEPINNLoss(
        alpha_func=problem['alpha_func'],
        f_func=problem['f_func'],
        level_set_func=problem['level_set_func'],
        tau_pde=1.0,
        tau_bc=100.0,
        tau_interface_u=100.0,
        tau_interface_flux=100.0
    )
    
    # Create solver
    solver = AEPINNSolver(
        model=model,
        loss_function=loss_function,
        domain_bounds=problem['domain_bounds'],
        level_set_func=problem['level_set_func']
    )
    
    # Setup optimizer
    solver.setup_optimizer(
        optimizer_type=args.optimizer,
        learning_rate=args.learning_rate,
        scheduler_type='step'
    )
    
    # Train
    solver.train(
        epochs=args.epochs,
        n_interior=args.n_interior,
        n_boundary=args.n_boundary,
        n_interface=args.n_interface,
        boundary_func=problem['boundary_func'],
        jump_func=problem['jump_func'],
        flux_jump_func=problem['flux_jump_func'],
        print_every=1000,
        save_every=5000,
        save_path=args.save_dir,
        use_adaptive_weights=True
    )
    
    print("Circular experiment completed!")


def run_custom_experiment(args):
    """Run a custom experiment from configuration file."""
    print(f"Running custom experiment from {args.config}")
    
    # Load configuration
    with open(args.config, 'r') as f:
        config = json.load(f)
    
    # Create problem from config
    problem = create_custom_problem_from_config(config.get('problem', {}))
    
    # Create model from config
    model_config = config.get('model', {})
    
    fcnn_config = model_config.get('fcnn', {})
    fcnn = create_fcnn({
        'input_dim': 2,
        'output_dim': 1,
        'hidden_layers': fcnn_config.get('layers', [50, 50, 50]),
        'activation': fcnn_config.get('activation', 'tanh')
    })
    
    ia_config = model_config.get('ia_nn', {})
    ia_nn = IANN(
        input_dim=2,
        hidden_dim=ia_config.get('hidden_dim', 40),
        num_attention_layers=ia_config.get('attention_layers', 3),
        output_dim=1,
        level_set_func=problem['level_set_func'],
        activation=ia_config.get('activation', 'tanh')
    )
    
    model = CompositeAEPINN(fcnn, ia_nn)
    
    # Training configuration
    training_config = config.get('training', {})
    
    # Create and run training
    loss_function = AEPINNLoss(
        alpha_func=problem['alpha_func'],
        f_func=problem['f_func'],
        level_set_func=problem['level_set_func'],
        **training_config.get('loss_weights', {})
    )
    
    solver = AEPINNSolver(
        model=model,
        loss_function=loss_function,
        domain_bounds=problem['domain_bounds'],
        level_set_func=problem['level_set_func']
    )
    
    solver.setup_optimizer(**training_config.get('optimizer', {}))
    solver.train(**training_config.get('train_params', {}))
    
    print("Custom experiment completed!")


def main():
    """Main function with command-line interface."""
    parser = argparse.ArgumentParser(
        description="AE-PINN: Attention-Enhanced Physics-Informed Neural Networks",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    
    # Experiment selection
    parser.add_argument('--experiment', type=str, default='star_shaped',
                        choices=['star_shaped', 'circular', 'custom'],
                        help='Type of experiment to run')
    
    parser.add_argument('--config', type=str, default=None,
                        help='Configuration file for custom experiments')
    
    # Problem parameters
    parser.add_argument('--r0', type=float, default=0.3,
                        help='Interface radius (for circular) or base radius (for star)')
    
    parser.add_argument('--alpha_in', type=float, default=1.0,
                        help='Diffusion coefficient inside interface')
    
    parser.add_argument('--alpha_out', type=float, default=5.0,
                        help='Diffusion coefficient outside interface')
    
    # Training parameters
    parser.add_argument('--epochs', type=int, default=10000,
                        help='Number of training epochs')
    
    parser.add_argument('--learning_rate', type=float, default=1e-3,
                        help='Learning rate')
    
    parser.add_argument('--optimizer', type=str, default='adam',
                        choices=['adam', 'lbfgs', 'sgd'],
                        help='Optimizer type')
    
    # Collocation points
    parser.add_argument('--n_interior', type=int, default=2000,
                        help='Number of interior collocation points')
    
    parser.add_argument('--n_boundary', type=int, default=400,
                        help='Number of boundary collocation points')
    
    parser.add_argument('--n_interface', type=int, default=400,
                        help='Number of interface collocation points')
    
    # Output
    parser.add_argument('--save_dir', type=str, default='results',
                        help='Directory to save results')
    
    parser.add_argument('--device', type=str, default=None,
                        help='Device to use (cpu, cuda, auto)')
    
    parser.add_argument('--seed', type=int, default=42,
                        help='Random seed for reproducibility')
    
    # Parse arguments
    args = parser.parse_args()
    
    # Set random seed
    torch.manual_seed(args.seed)
    np.random.seed(args.seed)
    
    # Set device
    if args.device is None or args.device == 'auto':
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    else:
        device = torch.device(args.device)
    
    print(f"Using device: {device}")
    
    # Initialize CUDA context if using GPU
    if device.type == 'cuda':
        initialize_cuda_context()
    
    # Create save directory
    if args.save_dir:
        os.makedirs(args.save_dir, exist_ok=True)
        print(f"Results will be saved to: {args.save_dir}")
    
    # Print configuration
    print("="*60)
    print("AE-PINN Configuration")
    print("="*60)
    for arg, value in vars(args).items():
        print(f"{arg:20s}: {value}")
    print("="*60)
    
    # Run experiment
    start_time = time.time()
    
    try:
        if args.experiment == 'star_shaped':
            run_star_shaped_experiment(args)
        elif args.experiment == 'circular':
            run_circular_experiment(args)
        elif args.experiment == 'custom':
            if args.config is None:
                raise ValueError("Configuration file required for custom experiments")
            run_custom_experiment(args)
        else:
            raise ValueError(f"Unknown experiment type: {args.experiment}")
    
    except KeyboardInterrupt:
        print("\nExperiment interrupted by user.")
        sys.exit(1)
    except Exception as e:
        print(f"\nError during experiment: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)
    
    total_time = time.time() - start_time
    print(f"\nExperiment completed in {total_time:.1f} seconds")


def create_example_config():
    """Create an example configuration file."""
    config = {
        "problem": {
            "type": "circular",
            "r0": 0.5,
            "alpha_in": 1.0,
            "alpha_out": 2.0
        },
        "model": {
            "fcnn": {
                "layers": [50, 50, 50],
                "activation": "tanh"
            },
            "ia_nn": {
                "hidden_dim": 40,
                "attention_layers": 3,
                "activation": "tanh"
            }
        },
        "training": {
            "optimizer": {
                "optimizer_type": "adam",
                "learning_rate": 1e-3,
                "scheduler_type": "step"
            },
            "train_params": {
                "epochs": 10000,
                "n_interior": 2000,
                "n_boundary": 400,
                "n_interface": 400,
                "print_every": 1000,
                "save_every": 5000,
                "use_adaptive_weights": True
            },
            "loss_weights": {
                "tau_pde": 1.0,
                "tau_bc": 100.0,
                "tau_interface_u": 100.0,
                "tau_interface_flux": 100.0
            }
        }
    }
    
    with open('example_config.json', 'w') as f:
        json.dump(config, f, indent=2)
    
    print("Example configuration saved to example_config.json")


if __name__ == "__main__":
    main()